# Script to produce a plot over various water deuteration measurements of astronomical objects.
# By Magnus V. Persson and John Tobin.
# 2022

# Requirements
"""
pandas
matplotlib
seaborn
numpy

"""

# should the plot be saved?
save = True

# Various imports 
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

import matplotlib.patches as mpatches
import matplotlib.transforms as transforms
from matplotlib.patches import Ellipse



# Set up base look of plots
sns.set(font_scale = 2)
sns.set_style("ticks")


# set tick width
plt.rcParams['ytick.major.size'] = 9
plt.rcParams['ytick.major.width'] = 3
plt.rcParams['ytick.minor.size'] = 6
plt.rcParams['ytick.minor.width'] = 2

plt.rcParams['lines.dash_joinstyle'] = 'round'        # default: miter|round|bevel
plt.rcParams['lines.dash_capstyle'] = 'round'         # default: butt|round|projecting
plt.rcParams['lines.solid_joinstyle'] = 'round'       # default: miter|round|bevel
plt.rcParams['lines.solid_capstyle'] = 'round'        # butt|round|default: projecting


# the base colors for the different measurements
c_earth = '#4c4c4c'
c_comets_occ = '#74409e'
c_comets_jfc = '#c16969'
c_protostars = '#405e9e'
c_cis = '#409e97f'
c_planets_jovian = '#9b873c'
c_moons = '#008000'
c_ism = '#009d82'
c_protosolar = '#0000ff'


# Define function to get various stats

def get_stats(data:pd.DataFrame()=None):
    """"""
    mean = np.mean(data)
    median = np.median(data)
    upper_quartile = np.percentile(data, 75)
    lower_quartile = np.percentile(data, 25)

    iqr = upper_quartile - lower_quartile
    upper_whisker = data[data<=upper_quartile+1.5*iqr].max()
    lower_whisker = data[data>=lower_quartile-1.5*iqr].min()
    return dict(mean=mean,
                median=median,
               uquart=upper_quartile,
               lquart=lower_quartile,
               uwhisker=upper_whisker,
               lwhisker=lower_whisker)


# Read the data file

df = pd.read_csv('ratios.csv', delimiter=';').dropna(how='all',axis=0)


# Now get on with the plotting directly


errlinewidth = 6
msize= 550

fig, ax = plt.subplots(figsize=(14,8))

ax_dh = ax.twinx()
ax_dh.set_yscale('log')


for _,s in ax.spines.items():
    s.set_linewidth(3)
    
for _,s in ax_dh.spines.items():
    s.set_linewidth(3)

############################### OC Comets
data_occ = df.query("target_type_main in ['Comet'] and target_type_sub_1 in ['Oort Cloud']")
data_occ = data_occ.assign(x=3+np.arange( len(data_occ)))
bplot = get_stats(data_occ.water_dh)

sns.lineplot(x=[data_occ.x.min()*0.55,data_occ.x.max()*1.015],
            y=[bplot['mean'],bplot['mean']],
            ax=ax,
             lw=6,
             alpha=0.5,
            color=c_comets_occ,
            zorder=-1,
            )

sns.scatterplot(data=data_occ, 
                x='x',y='water_dh', 
                s=msize, marker='o',
                ax=ax, 
                legend=False, 
                lw=0,
                color=c_comets_occ, 
                edgecolor=None
               )

ret = ax.errorbar(data_occ.x, data_occ.water_dh, yerr=data_occ.water_dh_std,
    ecolor=c_comets_occ, fmt=' ', zorder=-1, elinewidth=errlinewidth, solid_capstyle='round')
li,caps = ret.get_children()
caps.set_capstyle("round")


xdata = data_occ.x.values
ydata = data_occ.water_dh.values

trans = (fig.dpi_scale_trans +
         transforms.ScaledTranslation(np.mean(xdata)+0.2, np.mean(ydata)-0.0002, ax.transData))

# plot an ellipse around the point that is 150 x 130 points in diameter...
x_std, y_std = trans.transform([np.std(xdata)*0.4,np.std(ydata)*10000])/150
circle = mpatches.Ellipse((0, 0), x_std, y_std, angle=-35,
                          fill=True, lw=0, transform=trans, 
                          #color=c_comets_occ, zorder=-20, alpha=1)
                          color='#C1ADD2', zorder=-20, alpha=0.8)
ax.add_patch(circle)


ax.annotate('OCC', (np.mean(xdata),np.mean(ydata)+0.0025), ha='center')

############################### JF Comets

data_jfc = df.query("target_type_main in ['Comet'] and target_type_sub_1 in ['Jupiter Family']")
data_jfc = data_jfc.assign(x=4+data_occ.x.max()+np.arange( len(data_jfc)))

bplot = get_stats(data_jfc.water_dh)

sns.lineplot(x=[data_jfc.x.min()*0.94,data_jfc.x.max()*1.05],
            y=[bplot['mean'],bplot['mean']],
            ax=ax,
             lw=6,
             alpha=0.5,
            color=c_comets_jfc,
            zorder=-1,
            )

sns.scatterplot(data=data_jfc, 
                x='x',y='water_dh', 
                s=msize, marker='o',
                ax=ax, 
                legend=False, 
                lw=0,
                color=c_comets_jfc, 
                edgecolor=None
               )

ret = ax.errorbar(data_jfc.x, data_jfc.water_dh, yerr=data_jfc.water_dh_std,
    ecolor=c_comets_jfc, fmt=' ', zorder=-1, elinewidth=errlinewidth, solid_capstyle='round')
li,caps = ret.get_children()
caps.set_capstyle("round")


xdata = data_jfc.x.values
ydata = data_jfc.water_dh.values

trans = (fig.dpi_scale_trans +
         transforms.ScaledTranslation(np.mean(xdata), np.mean(ydata), ax.transData))

# plot an ellipse around the point that is 150 x 130 points in diameter...
x_std, y_std = trans.transform([np.std(xdata)*5,np.std(ydata)])/300
circle = mpatches.Ellipse((0, 0), x_std*0.4, y_std*2.1, angle=0,
                          fill=True, lw=0, transform=trans, 
                          #color=c_comets_jfc, zorder=-20, alpha=0.4)
                          color='#DFBDBD', zorder=-20, alpha=0.8)

ax.add_patch(circle)

ax.annotate('JFC', (np.mean(xdata),np.mean(ydata)+0.0015), ha='center')

ax.annotate('67P', (xdata[2]-0.6,ydata[2]-0.0001), ha='right', size=16)


############################# Class I
data_c1 = df.query("target_type_main in ['Protostar'] and target_type_sub_2 in ['Class I']")
data_c1 = data_c1.assign(x=4+data_jfc.x.max()+np.arange( len(data_c1)))
sns.scatterplot(data=data_c1, 
                x='x',y='water_dh',  
                s=400, 
                marker='s',
                ax=ax, 
                legend=False,
                color='#CCC348',
                edgecolor=None,
                )

ret = ax.errorbar(data_c1.x, data_c1.water_dh, yerr=data_c1.water_dh_std,
    ecolor='#CCC348', fmt=' ', zorder=-1, elinewidth=errlinewidth, solid_capstyle='round')

li,caps = ret.get_children()
caps.set_capstyle("round")

xdata = data_c1.x.values
ydata = data_c1.water_dh.values

trans = (fig.dpi_scale_trans +
         transforms.ScaledTranslation(np.mean(xdata), np.mean(ydata), ax.transData))

# plot an ellipse around the point that is 150 x 130 points in diameter...
x_std, y_std = trans.transform([1,5])/(150*6)
circle = mpatches.Ellipse((0, 0), x_std, y_std, angle=0,
                          fill=True, lw=0, transform=trans, 
                          #color='#CCC348',#c_protostars, 
                          color='#E3E0B0',#c_protostars, 
                          zorder=-20, alpha=0.8)
ax.add_patch(circle)

#ax.annotate('V883 Ori\nClass I Disk', (np.mean(xdata),np.mean(ydata)+0.0035), ha='center')
ax.annotate('Young proto-\nplanetary disk', (np.mean(xdata),np.mean(ydata)+0.0035), ha='center')

ax.annotate('V883Ori', (xdata[0]+0.6,ydata[0]+0.0005), ha='left', size=16)


############################# Class 0
data_c0 = df.query("target_type_main in ['Protostar'] and target_type_sub_2 in ['Class 0']")
data_c0 = data_c0.assign(x=4+data_c1.x.max()+np.arange( len(data_c0)))

bplot = get_stats(data_c0.water_dh)

sns.lineplot(x=[data_c0.x.min()*1.02,data_c0.x.max()*0.995],
            y=[bplot['mean'],bplot['mean']],
            ax=ax,
             lw=6,
             alpha=0.5,
            color=c_protostars,
            zorder=-1,
            )
sns.scatterplot(data=data_c0, 
                x='x',
                y='water_dh',  
                s=msize, 
                marker='o',
                color=c_protostars,
               ax=ax, legend=False,
               edgecolor=None)

ret = ax.errorbar(data_c0.x, data_c0.water_dh, yerr=data_c0.water_dh_std,
    ecolor='#3E5B99', fmt=' ', zorder=-1, elinewidth=errlinewidth, solid_capstyle='round')
li,caps = ret.get_children()
caps.set_capstyle("round")

xdata = data_c0.x.values
ydata = data_c0.water_dh.values

trans = (fig.dpi_scale_trans +
         transforms.ScaledTranslation(np.mean(xdata)+0.4, np.mean(ydata)*1.4, ax.transData))

# plot an ellipse around the point that is 150 x 130 points in diameter...
x_std, y_std = trans.transform([4,1])/(150*2.5)
circle = mpatches.Ellipse((-0.2, -0.46), x_std, y_std, angle=38,
                          fill=True, lw=0, transform=trans, 
                          #color=c_protostars, 
                          color='#ADB8D2', 
                          zorder=-20, alpha=0.8)
ax.add_patch(circle)

ax.annotate('Class 0', (np.mean(xdata),np.mean(ydata)+0.0025), ha='center')



#############################

ax.set_yscale('log')

ax.axhline(df.query("name == 'Earth'").iloc[0].water_dh, lw=4, zorder=-100, c=c_earth)
ax.annotate('Earth\'s oceans', (35.5,df.query("name == 'Earth'").iloc[0].water_dh*0.7), ha='center')


ax.axhline(df.query("name == 'Protosolar'").iloc[0].water_dh, lw=4, zorder=-100, c=c_protosolar)
ax.annotate('Protosolar', (36.5,df.query("name == 'Protosolar'").iloc[0].water_dh*1.15), ha='center')



ax.axhline(df.query("name == 'Local ISM'").iloc[0].water_dh, lw=4, zorder=-100, c=c_ism)
ax.annotate('Local ISM', (36.5,df.query("name == 'Local ISM'").iloc[0].water_dh*0.65), ha='center')


ax.set_ylim(9e-6,2e-2)
ax_dh.set_ylim([i/2 for i in ax.get_ylim()])
ax.set_xlim(0,40)


#sns.despine( bottom=True, top=True, ax=ax)
ax.tick_params(labelbottom=False)
ax.xaxis.set_ticks_position('none') 
ax.set_xlabel('')

#ax.tick_params(bottom=False,labelbottom=False)

ax.set_ylabel('log(HDO/H$_2$O)')

sns.despine( bottom=True, top=True, right=False, ax=ax_dh)
ax_dh.set_ylabel('log(D/H)')


ax.set_yticks([1e-2,1e-3,1e-4,1e-5])
ax.set_yticklabels(['-2','-3','-4','-5'])


ax_dh.set_yticks([1e-2,1e-3,1e-4,1e-5])
ax_dh.set_yticklabels(['-2','-3','-4','-5'])
plt.tight_layout()

# Finally save the figure
if save:
    plt.savefig('dh_ratio_lines.pdf')
    plt.savefig('dh_ratio_lines.eps')
else:
    plt.show()
